#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghuacheng <wanghuacheng@wanghuacheng-PC>
#
# Distributed under terms of the MIT license.

"""

"""
import sys
import os
import argparse
import baostock as bs
import pandas as pd
import talib as ta
import matplotlib.pyplot as plt
from talib.abstract import *
import datetime

date_range = [5, 13, 21, 34, 55, 89, 144, 233, 377]
len_date = len(date_range)

data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
merge_dir = os.path.join(data_dir, "history_k/merge/")
merge_dates = os.listdir(merge_dir)
md = max(merge_dates)
merge_latest_dir = os.path.join(merge_dir, md)
code_names = os.listdir(merge_latest_dir)
#print(codes)
#sys.exit(0)


today = datetime.datetime.today()
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-c', '--code', type=str, default=None, help='code name')
parser.add_argument('-a', '--action', type=str, default="st1", help='code name')

ARGS = parser.parse_args()

today_str = today.strftime("%Y-%m-%d")


def gold_cross(last_diff, today_diff):
    return last_diff < 0 and today_diff >= 0

def dead_cross(last_diff, today_diff):
    return last_diff > 0 and today_diff <= 0

def get_data(code, startdate, enddate):
    rs = bs.query_history_k_data(code,
                                 "date,code,close,low,high,open,tradeStatus,peTTM,pbMRQ,psTTM,pcfNcfTTM",
                                 start_date=startdate, end_date=enddate,
                                 frequency="d", adjustflag="3")
    df = rs.get_data()
    # print(df.tail())
    return df

def is_top(df, index):
    close = df.at[index, 'close']
    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close >= df.at[index + i, 'close'] and close >= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break
    return dflag

def is_bottom(df, index):
    close = df.at[index, 'close']
    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close <= df.at[index + i, 'close'] and close <= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break
    return dflag


def find_st1(df):
    print(df.head())
    N = df.shape[0]
    if N < 500:
        return None

    features = []

    # close features
    sma_close_diff_names = []
    for d in date_range:
        diff_name = 'SMA_DIFF_' + str(d)
        sma_name = 'SMA_' + str(d)
        sma_close_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d)
        df[diff_name] = (df['close'] - df[sma_name]) / df[sma_name]


    sma_diff_names = []
    for i in range(0, len_date - 1):
        for j in range(i + 1, len_date):
            diff_name = "SMA_" + str(date_range[i]) + "_" + str(date_range[j]) + "_DIFF"
            sma_diff_names.append(diff_name)
            df[diff_name] = (df['SMA_' + str(date_range[i])] - df['SMA_' + str(date_range[j])]) / df['SMA_' + str(date_range[j])]

    sma_fes = ['SMA_' + str(d) for d in date_range]

    features = sma_close_diff_names + sma_diff_names
    print(features)


    # val features .............
    val_features = ["peTTM", "pbMRQ", "psTTM", "pcfNcfTTM"]
    val_diff_features = []
    for vf in val_features:
        df[vf] = df[vf].astype('float')
        sma_name = vf + '_377_SMA'
        df[sma_name] = SMA(df, timeperiod = 377, price = vf)
        # vf_mean = df[vf].mean()
        # print(vf_mean)
        vf_diff_name = sma_name + "_diff"
        val_diff_features.append(vf_diff_name)
        df[vf_diff_name] = (df[vf] - df[sma_name])/ df[sma_name]


    # volume
    sma_vol_diff_names = []
    for d in date_range:
        diff_name = 'VOL_SMA_DIFF_' + str(d)
        sma_name = 'VOL_SMA_' + str(d)
        sma_vol_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d, price = 'volume')
        df[diff_name] = (df['volume'] - df[sma_name]) / df[sma_name]




    features = features + val_diff_features + sma_vol_diff_names

    for index, row in df.iterrows():
        if index < 500:
            continue
        if index > N - 10:
            break

        #print(row.date)
        dflag = is_bottom(df, index)
        tflag = is_top(df, index)
        if dflag:
            target = 0
            # print("BOTTOM:", row.date)
        if tflag:
            target = 1
            # print("TOP:", row.date)

        if dflag or tflag:
            # print(df.loc[index, features], target)
            fvals = (df.loc[index, features].values)
            yield fvals, target

def find_pe(df):
    #peMean = df['peTTM'].mean(skipna = True)
    pel = df.peTTM.astype('float')
    pe = (pel.mean())
    today_pe = pel[df.shape[0] - 1]
    print(today_pe, pe)
    # for index, row in df.iterrows():
        # print(index, row.peTTM)
    #print(peMean)


    #N = df.shape[0]
    #print(peMean, df.loc[N-1, 'peTTM'])

def find_macd(df):
    DIF = 1
    DEA = 2
    #df = df[df['tradeStatus'] == '1']
    if df.shape[0] < 50:
        return None
    macd_df = MACD(df, fastperiod=12, slowperiod=26, signalperiod=9)
    # print(macd_df)
    df = pd.concat([df, macd_df], axis = 1)
    N = int(df.shape[0])
    # print(df.shape)
    # dif 上穿 dea 白上穿黄
    # dif 下穿 dea 白下穿黄
    last_gold = None
    last_dead = None
    late_hist = df.at[N-1, 'macdhist']
    # print(df)
    for i in range(N - 1):
        day1_hist = df.at[i, 'macdhist']
        day2_hist = df.at[i + 1, 'macdhist']
        day2_dif =  df.at[i + 1, 'macd']
        dict1=df.loc[i + 1]
        dict2=df.loc[i]
        if gold_cross(day1_hist, day2_hist):
            # print(day2_dif > 0, "gold", dict2.to_json())
            last_gold = dict2['date']

        if dead_cross(day1_hist, day2_hist):
            # print(day2_dif > 0, "dead", dict2.to_json())
            last_dead = dict2['date']


    gddiff = get_date_diff(last_gold, last_dead)
    gdiff  = get_date_diff(today_str, last_gold)
    #print(last_gold, last_dead, late_hist)
    info = None
    if gdiff < 20 and gddiff > 0 and late_hist > 0:
        info = [last_dead, last_gold, gdiff]
    return info





def find_kdj(df):
    N = df.shape[0]
    matype = 0
    stoch = STOCH(df, fastk_period=9, slowk_matype=matype, slowk_period=3, slowd_period=3, slowd_matype=matype)
    df = pd.concat([df, stoch], axis=1)
    df['k_d'] = df['slowk'] - df['slowd']
    df['j'] = 3*df['slowk']- 2*df['slowd']

    last_dead = None
    last_gold = None
    last_kd = None
    for index, row in df.iterrows():
        if index < 50:
            continue
        last_kd = df.at[index - 1, 'k_d']
        today_kd = df.at[index, 'k_d']
        if gold_cross(last_kd, today_kd):
            last_gold = row.date
        elif dead_cross(last_kd, today_kd):
            last_dead = row.date

    gddiff = get_date_diff(last_gold, last_dead)
    gdiff  = get_date_diff(today_str, last_gold)
    info = None
    if gdiff < 20 and gddiff > 0 and last_kd > 0:
        info = [last_dead, last_gold, gdiff]

    return info



def find_rsi(df):
    N = df.shape[0]
    matype = 0
    rsis = []
    for d in [6, 12, 24]:
        dfn = 'rsi_' + str(d)
        df[dfn] = RSI(df, timeperiod = d)

        rsid = df.at[N-1 , dfn]
        rsis.append(rsid)

    score = 0
    # print(rsis)
    for index, r in enumerate(rsis):
        s = 0
        if r > 50 and r < 80:
            s = 1
        if r < 20:
            s = 1
        if r > 80:
            s = 1
        if r > 20 and r < 50:
            s = -1
        score += s * (index + 1)
    if score > 0:
        return str(score)
    else:
        return None

def find_bolling(df):
    N = df.shape[0]
    matype = 0
    rsis = []
    bands = BBANDS(df, timeperiod = 20, nbdevup = 2.0, nbdevdn = 2.0, matype = 0)
    df = pd.concat([df, bands], axis = 1)
    #upperband  middleband  lowerband
    # df.close - df.lowerband
    close = df.close.astype('float')
    df['BBP'] = (close - df.lowerband) / (df.upperband - df.lowerband)
    df['RSI'] = RSI(df, timeperiod = 14)
    x = df.loc[N-1, ['BBP', 'RSI']]
    print(x)


    score = 0
    print(df)
    if score > 0:
        return str(score)
    else:
        return None








def get_date_diff(d1, d2):
    dd1 = datetime.datetime.strptime(d1, '%Y-%m-%d')
    dd2 = datetime.datetime.strptime(d2, '%Y-%m-%d')
    return (dd1 - dd2).days

if __name__ == '__main__':
    # print(Function('stoch').info)
    login_result = bs.login(user_id='anonymous', password='123456')
    startdate = '2020-01-01'
    enddate = '2021-03-30'
    #code = 'sz.000921'
    code = ARGS.code
    action = ARGS.action
    fun = eval("find_"+ action)
    if code:
        df = get_data(code, startdate, enddate)
        info = fun(df)
        print(info)
    else:
        code_df = bs.query_hs300_stocks()
        all_df = code_df.get_data()
        outfile = action+"_gold.csv"

        with open(outfile, "w") as fw:
            for cn in code_names:
                code = cn[:9]
                code_name = cn[10:-4]
                cnfile = os.path.join(merge_latest_dir, cn)
                #print(cnfile)
                df = pd.read_csv(cnfile)
                info = fun(df)
                if info:
                    fw.write(code +","+ code_name +"," + ",".join(map(str, info)) + "\n")


        #with open(outfile, "w") as fw:
        #    for index,row in all_df.iterrows():
        #        df = get_data(row.code, startdate, enddate)
        #        info = fun(df)
        #        if info:
        #            fw.write(row.code +","+ row.code_name +"," + ",".join(map(str, info)) + "\n")
    bs.logout()
